from hep.aida import *

from java.awt import Color
from java.awt import Font
from java.lang import *
from java.util import *
from java.io import *

from jhplot.jadraw import *
from jhplot  import *

from fr.in2p3.ipnl.aidautilities import *
from fr.in2p3.ipnl.oflap.scint import *
from jscinttools..utilities.jhplotutilities import *

Locale.setDefault(Locale.ENGLISH)
System.getProperties().setProperty("hep.aida.BatchAnalysisFactory","hep.aida.ref.BatchAnalysisFactory")
af = IAnalysisFactory.create();

showFigures = True
writeFigures = not showFigures

""" Define colors and fonts """
figType = [".png", ".eps"]
colors = [Color.blue, Color.red, Color(0, 170, 0), Color.magenta, Color.black, Color.yellow]
axisLabelFont = Font(Font.SERIF, Font.PLAIN, 25)
titleFont = Font(Font.SERIF, Font.PLAIN, 25)
legendFont = Font(Font.SERIF, Font.PLAIN, 20)
ticFont = Font(Font.SERIF, Font.PLAIN, 20)
statBoxFont = Font(Font.SERIF, Font.PLAIN, 15)

""" Define the Path and files names """
aidaPathL = "/home/marc-antoine/Documents/BGO/QFin"
aidaFileL = "ScintillationAnalysis_BGO"
aidaPathQ = "/home/marc-antoine/Documents/BGO/QFin/QueensData/queenstools"
aidaFileQ = "ScintillationAnalysis_BGO_QFile"
outPath = "/home/marc-antoine/Dropbox/BGO"
afrL = AidaFileReader(aidaPathL, aidaFileL)
afrQ = AidaFileReader(aidaPathQ, aidaFileQ)

""" Create the DPS with the 2 initials DPS and do a fit """
mainDecay = ScintUtil.fusionDPS(afrL.getDPS("/Kinetic/Decay_Comp#2"),afrQ.getDPS("/Kinetic/Decay_Q_Comp#3"))
mainDecay = ScintUtil.removeNaN(mainDecay)

cutDPS = af.createDataPointSetFactory(af.createTreeFactory().create()).create(mainDecay.title(), 2);

j=0
for i in range(mainDecay.size()):
	if mainDecay.point(i).coordinate(0).value() < 100.0:
		cutDPS.addPoint();
		cutDPS.point(j).coordinate(0).setValue(mainDecay.point(i).coordinate(0).value());
		cutDPS.point(j).coordinate(0).setErrorMinus(mainDecay.point(i).coordinate(0).errorMinus());
		cutDPS.point(j).coordinate(0).setErrorPlus(mainDecay.point(i).coordinate(0).errorPlus());
		cutDPS.point(j).coordinate(1).setValue(mainDecay.point(i).coordinate(1).value());
		cutDPS.point(j).coordinate(1).setErrorMinus(mainDecay.point(i).coordinate(1).errorMinus());
		cutDPS.point(j).coordinate(1).setErrorPlus(mainDecay.point(i).coordinate(1).errorPlus());
		j+=1
			

#mainDecay = cutDPS

#kinParam = [350.0, 0.1, 100.0]
kinParam = [350.0, 0.1, 150.0]
#kinParam = [356.04, 9.58, 66.4, 0.0004, 3271.0]

doFit = False
fixShift = False

theResult = 0
linFunc = 0
if (len(kinParam) == 3):
		
	linFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("Afterglow Fit", 1, "1.0/((1.0/kag1_us)+(1.0/kag2_us)*exp(-1.0*dEag_K/x[0]))", "kag1_us,kag2_us,dEag_K", "Afterglow Fit")

	linFunc.setParameter("kag1_us", kinParam[0])
	linFunc.setParameter("kag2_us", kinParam[1])
	linFunc.setParameter("dEag_K", kinParam[2])

	if (fixShift):
		fitter.fitParameterSettings("kag1_us").setFixed(true)
		fitter.fitParameterSettings("kag2_us").setFixed(true)

		
elif (len(kinParam) == 5):
	linFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("Afterglow Fit", 1, "1.0/((1.0/kag1_us)+(1.0/kag2_us)*exp(-1.0*dEag_K/x[0])+(1.0+exp(-1.0*dEag_K/x[0]))*(1.0/capKag_us)*exp(-1.0*deltaEag_K/x[0]))", "kag1_us,kag2_us,dEag_K,capKag_us,deltaEag_K", "Afterglow Fit")

	linFunc.setParameter("kag1_us", kinParam[0])
	linFunc.setParameter("kag2_us", kinParam[1])
	linFunc.setParameter("dEag_K", kinParam[2])
	linFunc.setParameter("capKag_us", kinParam[3])
	linFunc.setParameter("deltaEag_K", kinParam[4])

	if (fixShift):
		fitter.fitParameterSettings("kag1_us").setFixed(true)
		fitter.fitParameterSettings("kag2_us").setFixed(true)

if doFit:
	theResult = fitter.fit(mainDecay, linFunc);
	print("\n================ Afterglow Fit ================")
	print(AidaUtil.printResultsOfFit(theResult))
	
	linFunc = theFitResult.fittedFunction()


#theFitResult = ScintDataFitter.doAfterglowFit(mainDecay, kinParam)
"""
###############################################################
# Plotting Part
###############################################################
"""
""" Create the plottable JHep objects from AIDA IObjects """
""" The DataPointSet """
mainP1D = P1D(mainDecay)
mainP1D.setTitle("Main Component")
mainP1D.setStyle("p")
mainP1D.setSymbol(2)
mainP1D.setSymbolSize(12)
mainP1D.setColor(colors[3])
mainP1D.setPenWidth(2)
mainP1D.setErrY(True)
mainP1D.setErrX(False)
mainP1D.setErrColor(mainP1D.getColor())

mainP1DL = P1D(afrL.getDPS("/Kinetic/Decay_Comp#2"))
mainP1DL.setTitle("Crystal 1")
mainP1DL.setStyle("p")
mainP1DL.setSymbol(6)
mainP1DL.setSymbolSize(12)
mainP1DL.setColor(colors[0])
mainP1DL.setPenWidth(2)
mainP1DL.setErrY(True)
mainP1DL.setErrX(False)
mainP1DL.setErrColor(mainP1DL.getColor())

mainP1DQ = P1D(afrQ.getDPS("/Kinetic/Decay_Q_Comp#3"))
mainP1DQ.setTitle("Crystal 2")
mainP1DQ.setStyle("p")
mainP1DQ.setSymbol(8)
mainP1DQ.setSymbolSize(16)
mainP1DQ.setColor(colors[1])
mainP1DQ.setPenWidth(2)
mainP1DQ.setErrY(True)
mainP1DQ.setErrX(False)
mainP1DQ.setErrColor(mainP1DQ.getColor())

""" The gamma fit fuction """
mainF1D = F1D(linFunc, 0.1, 1000)
mainF1D.setTitle("Fit &gamma;")
mainF1D.setPoints(10000)
mainF1D.setColor(colors[2])
mainF1D.setPenWidth(2)


mainP1DKeyL = JaKey(mainP1DL.getSymbol(), mainP1DL.getTitle(), 0.20, 0.65, "NDC")
mainP1DKeyL.setColor(mainP1DL.getColor())
mainP1DKeyL.setKeyColor(mainP1DL.getColor())
mainP1DKeyL.setPenWidth(2)
mainP1DKeyL.setFont(legendFont)

mainP1DKeyQ = JaKey(mainP1DQ.getSymbol(), mainP1DQ.getTitle(), 0.20, 0.70, "NDC")
mainP1DKeyQ.setColor(mainP1DQ.getColor())
mainP1DKeyQ.setKeyColor(mainP1DQ.getColor())
mainP1DKeyQ.setPenWidth(2)
mainP1DKeyQ.setFont(legendFont)

mainF1DKey = JaKey(20, mainF1D.getTitle(), 0.20, 0.75, "NDC")
mainF1DKey.setColor(mainF1D.getColor())
mainF1DKey.setKeyColor(mainF1D.getColor())
mainF1DKey.setPenWidth(2)
mainF1DKey.setFont(legendFont)


""" Create the Plotter """
mainPlotter = HPlotJa("Fusion Kinetics", 600, 600, 1, 1, True)
mainPlotter.setGTitle("Decay Time in BGO", titleFont)
mainPlotter.setAntiAlias(False)
if showFigures:
	mainPlotter.visible(True)
else:
	mainPlotter.visible(False)
mainPlotter.setPad(1, 1, 0.125, 0.125, 0.75, 0.75)
mainPlotter.setRange(1, 1000, 0.1, 1000)
mainPlotter.setLogScale(0, True)
mainPlotter.setLogScale(1, True)
mainPlotter.setTicFont(ticFont)
mainPlotter.setShowStatBox(False)
mainPlotter.showKey(False)
mainPlotter.setNameX("Temperature (K)")
mainPlotter.setNameY("Decay Time (&micro;s)")

""" Plots all the JHep objects """
mainPlotter.draw(mainP1D)
mainPlotter.draw(mainF1D)
#mainPlotter.draw(mainP1DL)
#mainPlotter.draw(mainP1DQ)
#mainPlotter.draw(mainP1DKeyL)
#mainPlotter.draw(mainP1DKeyQ)
#mainPlotter.draw(mainF1DKey)

""" Modifies the JaAxes and the axes labels """
labX = mainPlotter.getLabelX()[0][0]
labX.setFont(axisLabelFont)
HPlotJaUtil.setLabelLocationAtCenter(0, labX, mainPlotter)
labY = mainPlotter.getLabelY()[0][0]
labY.setFont(axisLabelFont)
HPlotJaUtil.setLabelLocationAtCenter(1, labY, mainPlotter)

theAxes = mainPlotter.getPad()
theAxes.setLabelSpace(0, 0.0)
theAxes.setLabelSpace(1, -0.025)
HPlotJaUtil.setTickLabelOneToPower(theAxes);

mainPlotter.update()

""" Saves or not the plot """
if writeFigures:
	for i in range(len(figType)):
		outFileName = "FusKin3lvlAg" + figType[i]
		mainPlotter.export(outPath + File.separator + outFileName)
		print("Writting File : " + outPath + File.separator + outFileName)
		
mainPlotter.update()
print("That's all folks !")

